from src.graph.builder import *
import uuid
import json
from langchain_core.messages import HumanMessage


def convert_messages_for_state(messages):
    """将 LangChain 消息对象转换为字典格式"""
    formatted = []
    for msg in messages:
        if isinstance(msg, dict):
            formatted.append(msg)
        else:
            # 从 LangChain 消息对象提取信息
            if hasattr(msg, 'type') and hasattr(msg, 'content'):
                role_map = {
                    'human': 'user',
                    'ai': 'assistant',
                    'system': 'system'
                }
                role = role_map.get(getattr(msg, 'type', 'assistant'), 'assistant')
                formatted.append({"role": role, "content": msg.content})
            else:
                formatted.append({"role": "assistant", "content": str(msg)})
    return formatted


def state_to_dict(state):
    """将State对象转换为可序列化的字典"""
    state_dict = {}
    for key, value in state.items():
        if key == "messages":
            # 转换消息为字典格式
            state_dict[key] = convert_messages_for_state(value)
        else:
            # 其他字段直接保留
            state_dict[key] = value
    return state_dict


def main():
    print("=== AI 对话系统 ===")
    print("输入 'quit' 结束对话，'new' 开始新对话")
    print("-" * 30)

    conversation_id = f"conv_{uuid.uuid4()}"
    messages = []

    while True:
        user_input = input("\n👤 你: ").strip()

        if user_input.lower() in ['quit', 'exit']:
            print("👋 再见！")
            break
        elif user_input.lower() == 'new':
            conversation_id = f"conv_{uuid.uuid4()}"
            messages = []
            print("🔄 新对话已开始")
            continue
        elif not user_input:
            continue

        # 添加用户消息
        messages.append(HumanMessage(content=user_input))

        # 转换消息格式
        formatted_messages = convert_messages_for_state(messages)

        # 创建状态
        state = State(
            messages=formatted_messages,
            locale="zh-CN",
            research_topic="General Conversation",
            enable_background_investigation=True,
            background_investigation_results=""
        )

        # 调用 graph
        try:
            result = graph.invoke(
                state,
                config={"configurable": {"thread_id": conversation_id}}
            )

            # 将state转换为可序列化的字典并打印
            state_dict = state_to_dict(result)
            print("\n📊 当前State:")
            print(json.dumps(state_dict, indent=2, ensure_ascii=False))
            print()

            # 更新消息历史
            messages = result["messages"]

        except Exception as e:
            print(f"❌ 错误: {e}")
            import traceback
            traceback.print_exc()


if __name__ == "__main__":
    main()